package com.yanghui.dtracing.core.opentracing;

import com.yanghui.dtracing.core.opentracing.propagation.Extractor;
import com.yanghui.dtracing.core.opentracing.propagation.Injector;
import com.yanghui.dtracing.core.opentracing.propagation.PropagationRegister;
import com.yanghui.dtracing.core.opentracing.propagation.TextMapJsonCodec;
import com.yanghui.dtracing.core.opentracing.reporter.Reporter;
import com.yanghui.dtracing.core.opentracing.reporter.ReporterManager;
import com.yanghui.dtracing.core.opentracing.sampler.Sampler;
import com.yanghui.dtracing.core.opentracing.sampler.SamplerProperties;
import com.yanghui.dtracing.core.opentracing.sampler.TracerPercentageBasedSampler;
import io.opentracing.*;
import io.opentracing.propagation.Format;
import io.opentracing.tag.Tag;
import io.opentracing.util.ThreadLocalScopeManager;

import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * @author nick.yang
 */
public class DTracer implements Tracer {

    private final ScopeManager scopeManager = new ThreadLocalScopeManager();

    private Reporter reporter;

    private TraceIdGenerate traceIdGenerate = new RuleTraceGenerate();

    private Sampler sampler;

    private PropagationRegister propagationRegister;

    private String serviceName;

    private DTracer(){}

    void report(DTracerSpan span){
        if(null == span){
            return;
        }
        if(!span.context().isSampled()){
            return;
        }
        if(null != reporter){
            this.reporter.report(span);
        }
    }

    @Override
    public ScopeManager scopeManager() {
        return this.scopeManager;
    }

    @Override
    public Span activeSpan() {
        return this.scopeManager.activeSpan();
    }

    @Override
    public Scope activateSpan(Span span) {
        Scope scope =  this.scopeManager.activate(span);
        ScopeContext.put(span.context().toSpanId(),scope);
        return scope;
    }

    @Override
    public SpanBuilder buildSpan(String operationName) {
        SpanBuilder spanBuilder = new SpanBuilder(operationName);
        spanBuilder.withStartTimestamp(System.currentTimeMillis());
        return spanBuilder;
    }

    @Override
    public <C> void inject(SpanContext spanContext, Format<C> format, C carrier) {
        Injector<C> injector = this.propagationRegister.getInjector(format);
        if(null == injector){
            throw new RuntimeException("Unsupported Injector of Format:" + format.toString());
        }
        injector.inject(spanContext,carrier);
    }

    @Override
    public <C> SpanContext extract(Format<C> format, C carrier) {
        Extractor<C> extractor = this.propagationRegister.getExtractor(format);
        if(null == extractor){
            throw new RuntimeException("Unsupported Extractor of Format:" + format.toString());
        }
        return extractor.extract(carrier);
    }

    @Override
    public void close() {
        this.reporter.close();
        this.sampler.close();
    }

    public static Builder builder(String serviceName){
        return new Builder(serviceName);
    }

    public static class Builder{

        private Reporter reporter;

        private Sampler sampler;

        private PropagationRegister propagationRegister;

        private String serviceName;

        private Builder(String serviceName){
            this.serviceName  =serviceName;
        }

        public Builder sampler(SamplerProperties samplerProperties){
            this.sampler = new TracerPercentageBasedSampler(samplerProperties);
            return this;
        }

        public Builder reporter(){
            this.reporter = ReporterManager.loadReporters();
            return this;
        }

        public Builder register(){
            this.propagationRegister = new PropagationRegister();
            TextMapJsonCodec textMapJsonCodec = new TextMapJsonCodec();
            propagationRegister.registerInjector(Format.Builtin.TEXT_MAP_INJECT, textMapJsonCodec);
            propagationRegister.registerExtractor(Format.Builtin.TEXT_MAP_EXTRACT,textMapJsonCodec);
            return this;
        }

        public DTracer build(){
            DTracer tracer = new DTracer();
            tracer.propagationRegister = propagationRegister;
            tracer.reporter = reporter;
            tracer.sampler = sampler;
            tracer.serviceName = serviceName;
            return tracer;
        }
    }

    public class SpanBuilder implements Tracer.SpanBuilder{

        private final String operationName;

        private String className;

        private long startTime;

        private List<SpanContextReference> referenceList = Collections.emptyList();

        private Map<String,Object> tags = new HashMap<>();

        public SpanBuilder(String operationName){
            this.operationName = operationName;
        }

        public SpanBuilder withClassName(String className){
            this.className = className;
            return this;
        }

        @Override
        public Tracer.SpanBuilder asChildOf(SpanContext parent) {
            if(null == parent){
                return this;
            }
            if(!(parent instanceof DTracerSpanContext)){
                return this;
            }
            SpanContextReference spanContextReference = new SpanContextReference(References.CHILD_OF,(DTracerSpanContext)parent);
            if(this.referenceList.isEmpty()){
                referenceList = Collections.singletonList(spanContextReference);
            }else {
                /**这里会抛出异常，目前只支持只一个父节点的情况 **/
                referenceList.add(spanContextReference);
            }
            return this;
        }

        @Override
        public Tracer.SpanBuilder asChildOf(Span parent) {
            return this.asChildOf(parent.context());
        }

        @Override
        public Tracer.SpanBuilder addReference(String referenceType, SpanContext referencedContext) {
            if(null == referencedContext){
                return this;
            }
            if (!(referencedContext instanceof DTracerSpanContext)) {
                return this;
            }
            if (!References.CHILD_OF.equals(referenceType)
                    && !References.FOLLOWS_FROM.equals(referenceType)) {
                return this;
            }
            SpanContextReference spanContextReference = new SpanContextReference(referenceType,(DTracerSpanContext)referencedContext);
            if(this.referenceList.isEmpty()){
                this.referenceList = Collections.singletonList(spanContextReference);
            }else {
                /**这里会抛出异常，目前只支持只一个父节点的情况 **/
                referenceList.add(spanContextReference);
            }
            return this;
        }

        @Override
        public Tracer.SpanBuilder ignoreActiveSpan() {
            //TODO 未实现
            return this;
        }

        @Override
        public Tracer.SpanBuilder withTag(String key, String value) {
            this.tags.put(key,value);
            return this;
        }

        @Override
        public Tracer.SpanBuilder withTag(String key, boolean value) {
            this.tags.put(key,value);
            return this;
        }

        @Override
        public Tracer.SpanBuilder withTag(String key, Number value) {
            this.tags.put(key,value);
            return this;
        }

        @Override
        public <T> Tracer.SpanBuilder withTag(Tag<T> tag, T value) {
            this.tags.put(tag.getKey(),value);
            return this;
        }

        @Override
        public Tracer.SpanBuilder withStartTimestamp(long microseconds) {
            this.startTime = microseconds;
            return this;
        }

        @Override
        public Span start() {
            DTracerSpanContext spanContext;
            Span parent = scopeManager.activeSpan();
            if(referenceList.isEmpty() && null != parent){
                this.asChildOf(parent);
            }
            /**无父节点，创建新的上下文**/
            if(referenceList.isEmpty()){
                spanContext = createNewContext();
            }else {
                spanContext = createChildContext();
            }
            DTracerSpan span = new DTracerSpan(DTracer.this,spanContext,serviceName,className,operationName,startTime,this.tags);
            /** 根span进行采样计算**/
            if(span.isRootSpan() && sampler != null){
                boolean isSampled = sampler.sample(span);
                spanContext.setSampled(isSampled);
            }
            return span;
        }

        private DTracerSpanContext createChildContext() {
            DTracerSpanContext parentContext = this.referenceList.get(0).getSpanContext();
            String spanId = parentContext.nextChildSpanId();
            DTracerSpanContext newContext = new DTracerSpanContext(parentContext.toTraceId(),
                    spanId,
                    parentContext.getSpanId(),
                    parentContext.isSampled());
            newContext.addBaggageItem(Baggage.ASYN_KEY,parentContext.getBaggageItem(Baggage.ASYN_KEY));
            String autoId = parentContext.getBaggageItem(Baggage.AUTO_ID_KEY);
            if(autoId != null && !"".equals(autoId)){
                newContext.addBaggageItem(Baggage.AUTO_ID_KEY,String.valueOf(ScopeContext.increment(Integer.valueOf(autoId))));
            }else {
                newContext.addBaggageItem(Baggage.AUTO_ID_KEY,String.valueOf(ScopeContext.increment(0)));
            }

            return newContext;
        }

        private DTracerSpanContext createNewContext(){
            DTracerSpanContext newContext = new DTracerSpanContext(traceIdGenerate.createTraceId());
            newContext.addBaggageItem(Baggage.ASYN_KEY,Baggage.ASYN_VALUE_NO);
            newContext.addBaggageItem(Baggage.AUTO_ID_KEY,String.valueOf(ScopeContext.increment(0)));
            return newContext;
        }
    }
}
